from PyQt6.QtWidgets import (QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
                             QPushButton, QLabel, QFileDialog, QTableWidget,
                             QTableWidgetItem, QMessageBox, QCheckBox, QApplication,
                             QHeaderView, QMenu, QTreeWidget, QTreeWidgetItem)
from PyQt6.QtCore import Qt
import pandas as pd
from utils.file_counter import FileCounter
from pathlib import Path
import os

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("文件统计工具")
        self.setMinimumSize(1000, 700)  # 增加窗口默认大小
        
        # 创建主窗口部件
        main_widget = QWidget()
        self.setCentralWidget(main_widget)
        
        # 创建主布局
        layout = QVBoxLayout()
        main_widget.setLayout(layout)
        
        # 创建顶部控制区域
        control_layout = QHBoxLayout()
        
        # 添加选择文件/文件夹按钮
        self.select_btn = QPushButton("选择文件/文件夹")
        self.select_btn.setContextMenuPolicy(Qt.ContextMenuPolicy.CustomContextMenu)
        self.select_btn.customContextMenuRequested.connect(self.show_select_menu)
        self.select_btn.clicked.connect(self.select_files)
        control_layout.addWidget(self.select_btn)
        
        # 添加统计按钮
        self.count_btn = QPushButton("开始统计")
        self.count_btn.clicked.connect(self.start_counting)
        control_layout.addWidget(self.count_btn)
        
        # 添加复制按钮
        self.copy_btn = QPushButton("复制结果")
        self.copy_btn.clicked.connect(self.copy_results)
        control_layout.addWidget(self.copy_btn)
        
        # 添加导出按钮
        self.export_btn = QPushButton("导出结果")
        self.export_btn.clicked.connect(self.export_results)
        control_layout.addWidget(self.export_btn)
        
        layout.addLayout(control_layout)
        
        # 创建树形视图
        self.tree = QTreeWidget()
        self.tree.setHeaderLabels(["文件路径", "文件数量", "代码行数", "总行数", "状态"])
        self.tree.setColumnWidth(0, 400)  # 设置第一列宽度
        self.tree.setStyleSheet("""
            QTreeWidget {
                gridline-color: #d0d0d0;
                background-color: white;
                alternate-background-color: #f6f6f6;
            }
            QTreeWidget::item {
                padding: 4px;
            }
            QHeaderView::section {
                background-color: #f0f0f0;
                padding: 4px;
                border: 1px solid #d0d0d0;
                font-weight: bold;
            }
        """)
        
        # 添加右键菜单
        self.tree.setContextMenuPolicy(Qt.ContextMenuPolicy.CustomContextMenu)
        self.tree.customContextMenuRequested.connect(self.show_tree_menu)
        
        self.tree.itemExpanded.connect(self.on_item_expanded)
        
        layout.addWidget(self.tree)
        
        # 创建状态栏
        self.statusBar().showMessage("就绪")
        
        # 初始化文件计数器
        self.file_counter = FileCounter()
        self.selected_files = set()  # 使用集合来存储文件路径，自动去重
        self.file_stats = {}  # 存储文件统计结果
        self.total_code_lines_label = QLabel("总代码行数：0")
        self.total_code_lines_label.setStyleSheet("padding-right: 12px; color: #333;")
        self.total_code_lines_label.setAlignment(Qt.AlignmentFlag.AlignRight | Qt.AlignmentFlag.AlignVCenter)
        self.total_code_lines_label.setCursor(Qt.CursorShape.PointingHandCursor)
        self.total_code_lines_label.mousePressEvent = self.copy_total_code_lines
        self.statusBar().addPermanentWidget(self.total_code_lines_label)

        self.total_lines_label = QLabel("总行数：0")
        self.total_lines_label.setStyleSheet("padding-right: 12px; color: #333;")
        self.total_lines_label.setAlignment(Qt.AlignmentFlag.AlignRight | Qt.AlignmentFlag.AlignVCenter)
        self.total_lines_label.setCursor(Qt.CursorShape.PointingHandCursor)
        self.total_lines_label.mousePressEvent = self.copy_total_lines
        self.statusBar().addPermanentWidget(self.total_lines_label)
        
    def show_select_menu(self, pos):
        menu = QMenu(self)
        select_files_action = menu.addAction("选择文件")
        select_dir_action = menu.addAction("选择文件夹")
        
        action = menu.exec(self.select_btn.mapToGlobal(pos))
        if action == select_files_action:
            self.select_files()
        elif action == select_dir_action:
            self.select_directory()
            
    def show_tree_menu(self, pos):
        menu = QMenu(self)
        expand_action = menu.addAction("展开")
        collapse_action = menu.addAction("折叠")
        menu.addSeparator()
        remove_action = menu.addAction("移除选中项")
        clear_action = menu.addAction("清空列表")
        
        item = self.tree.itemAt(pos)
        if item:
            action = menu.exec(self.tree.mapToGlobal(pos))
            if action == expand_action:
                self.tree.expandItem(item)
            elif action == collapse_action:
                self.tree.collapseItem(item)
            elif action == remove_action:
                self.remove_selected_items()
            elif action == clear_action:
                self.clear_tree()
        else:
            action = menu.exec(self.tree.mapToGlobal(pos))
            if action == remove_action:
                self.remove_selected_items()
            elif action == clear_action:
                self.clear_tree()
            
    def remove_selected_items(self):
        selected_items = self.tree.selectedItems()
        if not selected_items:
            return
            
        for item in selected_items:
            path = self.get_full_path(item)
            if path in self.selected_files:
                self.selected_files.remove(path)
            parent = item.parent()
            if parent:
                parent.removeChild(item)
            else:
                self.tree.takeTopLevelItem(self.tree.indexOfTopLevelItem(item))
            
        self.update_status()
        
    def clear_tree(self):
        self.tree.clear()
        self.selected_files.clear()
        self.file_stats.clear()
        self.update_status()
        
    def select_files(self):
        files, _ = QFileDialog.getOpenFileNames(
            self,
            "选择文件",
            "",
            "所有文件 (*.*)"
        )
        if files:
            self.add_files(files)
            
    def select_directory(self):
        dir_path = QFileDialog.getExistingDirectory(
            self,
            "选择文件夹",
            ""
        )
        if dir_path:
            self.add_files([dir_path])
            
    def add_files(self, paths):
        """添加文件到列表，自动去重并排序"""
        for path in paths:
            self.selected_files.add(path)
        self.selected_files = set(sorted(self.selected_files, key=lambda x: x.lower()))
        self.refresh_tree()
        self.update_status()

    def refresh_tree(self):
        self.tree.clear()
        for path in sorted(self.selected_files, key=lambda x: x.lower()):
            self.add_item_to_tree(path)

    def add_item_to_tree(self, path):
        item = QTreeWidgetItem()
        item.setText(0, os.path.basename(path))
        item.setText(1, "待统计")
        item.setText(2, "待统计")
        item.setText(3, "待统计")
        item.setText(4, "未统计")
        if os.path.isdir(path):
            item.setChildIndicatorPolicy(QTreeWidgetItem.ChildIndicatorPolicy.ShowIndicator)
            item.setIcon(0, self.style().standardIcon(self.style().StandardPixmap.SP_DirIcon))
        else:
            item.setIcon(0, self.style().standardIcon(self.style().StandardPixmap.SP_FileIcon))
        self.tree.addTopLevelItem(item)

    def update_status(self):
        total_files = len(self.selected_files)
        self.statusBar().showMessage(f"已选择 {total_files} 个文件/文件夹")
        self.total_lines_label.setText("总行数：0")
        self.total_code_lines_label.setText("总代码行数：0")

    def start_counting(self):
        if not self.selected_files:
            QMessageBox.warning(self, "警告", "请先选择文件或文件夹")
            return
        try:
            QApplication.setOverrideCursor(Qt.CursorShape.WaitCursor)
            self.statusBar().showMessage("正在统计，请稍候...")
            QApplication.processEvents()
            self.file_stats.clear()
            self.tree.clear()
            for path in sorted(self.selected_files, key=lambda x: x.lower()):
                self.add_item_to_tree(path)
            results = self.file_counter.count_files(
                list(self.selected_files),
                recursive=True
            )
            self.file_stats = {path: (file_count, line_count, total_lines)
                             for path, file_count, line_count, total_lines in results}
            self.display_results(results)
            self.statusBar().showMessage("统计完成")
            # 统计总行数和总代码行数（只统计顶层选中项）
            total_lines = sum(self.file_stats.get(path, (0, 0, 0))[2] for path in self.selected_files)
            total_code_lines = sum(self.file_stats.get(path, (0, 0, 0))[1] for path in self.selected_files)
            self.total_lines_label.setText(f"总行数：{total_lines}")
            self.total_code_lines_label.setText(f"总代码行数：{total_code_lines}")
        except Exception as e:
            QMessageBox.critical(self, "错误", f"统计过程中发生错误：{str(e)}")
        finally:
            QApplication.restoreOverrideCursor()

    def display_results(self, results):
        for path, file_count, line_count, total_lines in results:
            items = self.tree.findItems(os.path.basename(path), Qt.MatchFlag.MatchExactly, 0)
            if items:
                item = items[0]
                item.setText(1, str(file_count))
                item.setText(2, str(line_count))
                item.setText(3, str(total_lines))
                item.setText(4, "已统计")
                if os.path.isdir(path):
                    self.update_folder_contents(item, path)
                    self.tree.expandItem(item)

    def update_folder_contents(self, parent_item, folder_path):
        try:
            contents = self.file_counter.get_directory_contents(
                folder_path,
                recursive=True
            )
            path_to_item = {folder_path: parent_item}
            contents.sort(key=lambda x: x[5])
            for path, is_dir, file_count, line_count, total_lines, rel_path in contents:
                parent_path = os.path.dirname(path)
                if parent_path not in path_to_item:
                    continue
                parent_item = path_to_item[parent_path]
                child_item = None
                for i in range(parent_item.childCount()):
                    if parent_item.child(i).text(0) == os.path.basename(path):
                        child_item = parent_item.child(i)
                        break
                if not child_item:
                    child_item = QTreeWidgetItem()
                    child_item.setText(0, os.path.basename(path))
                    parent_item.addChild(child_item)
                if is_dir:
                    child_item.setIcon(0, self.style().standardIcon(self.style().StandardPixmap.SP_DirIcon))
                    child_item.setChildIndicatorPolicy(QTreeWidgetItem.ChildIndicatorPolicy.ShowIndicator)
                    path_to_item[path] = child_item
                else:
                    child_item.setIcon(0, self.style().standardIcon(self.style().StandardPixmap.SP_FileIcon))
                child_item.setText(1, str(file_count))
                child_item.setText(2, str(line_count))
                child_item.setText(3, str(total_lines))
                child_item.setText(4, "已统计")
        except Exception as e:
            print(f"Error updating folder {folder_path}: {str(e)}")

    def copy_results(self):
        if self.tree.topLevelItemCount() == 0:
            QMessageBox.warning(self, "警告", "没有可复制的数据")
            return
        text = []
        headers = []
        for i in range(self.tree.columnCount()):
            headers.append(self.tree.headerItem().text(i))
        text.append("\t".join(headers))
        def add_item_to_text(item, level=0):
            row_data = []
            indent = "    " * level
            row_data.append(f"{indent}{item.text(0)}")
            row_data.append(item.text(1))
            row_data.append(item.text(2))
            row_data.append(item.text(3))
            row_data.append(item.text(4))
            text.append("\t".join(row_data))
            for i in range(item.childCount()):
                add_item_to_text(item.child(i), level + 1)
        for i in range(self.tree.topLevelItemCount()):
            add_item_to_text(self.tree.topLevelItem(i))
        clipboard = QApplication.clipboard()
        clipboard.setText("\n".join(text))
        self.statusBar().showMessage("结果已复制到剪贴板")
        
    def copy_total_lines(self, event):
        text = self.total_lines_label.text().replace("总行数：", "").strip()
        clipboard = QApplication.clipboard()
        clipboard.setText(text)
        self.statusBar().showMessage("总行数已复制到剪贴板")
        
    def copy_total_code_lines(self, event):
        text = self.total_code_lines_label.text().replace("总代码行数：", "").strip()
        clipboard = QApplication.clipboard()
        clipboard.setText(text)
        self.statusBar().showMessage("总代码行数已复制到剪贴板")
        
    def export_results(self):
        if self.tree.topLevelItemCount() == 0:
            QMessageBox.warning(self, "警告", "没有可导出的数据")
            return
        file_path, _ = QFileDialog.getSaveFileName(
            self,
            "导出结果",
            "",
            "Excel文件 (*.xlsx);;CSV文件 (*.csv)"
        )
        if not file_path:
            return
        try:
            data = []
            headers = [self.tree.headerItem().text(i) for i in range(self.tree.columnCount())]
            def add_item_to_data(item, level=0):
                row_data = []
                indent = "    " * level
                row_data.append(f"{indent}{item.text(0)}")
                row_data.append(item.text(1))
                row_data.append(item.text(2))
                row_data.append(item.text(3))
                row_data.append(item.text(4))
                data.append((row_data, item))
                for i in range(item.childCount()):
                    add_item_to_data(item.child(i), level + 1)
            for i in range(self.tree.topLevelItemCount()):
                add_item_to_data(self.tree.topLevelItem(i))
            df = pd.DataFrame([row for row, _ in data], columns=headers)
            if file_path.endswith('.xlsx'):
                df.to_excel(file_path, index=False)
                try:
                    from openpyxl import load_workbook
                    from openpyxl.styles import Font
                    wb = load_workbook(file_path)
                    ws = wb.active
                    for idx, (row, item) in enumerate(data, start=2):
                        if item.childIndicatorPolicy() == QTreeWidgetItem.ChildIndicatorPolicy.ShowIndicator:
                            for cell in ws[idx]:
                                cell.font = Font(color="FF0000")
                    wb.save(file_path)
                except Exception as e:
                    print(f"Excel着色失败: {e}")
            else:
                df.to_csv(file_path, index=False)
            self.statusBar().showMessage("导出成功")
        except Exception as e:
            QMessageBox.critical(self, "错误", f"导出过程中发生错误：{str(e)}")
            
    def on_item_expanded(self, item):
        """动态加载文件夹下一级内容（只显示结构，不统计）"""
        # 如果已经有子项则不再重复加载
        if item.childCount() > 0 and item.child(0).text(1) != "待统计":
            return
        # 清空原有子项（防止重复加载）
        item.takeChildren()
        # 获取完整路径
        path = self.get_full_path(item)
        if not path or not os.path.isdir(path):
            return
        try:
            with os.scandir(path) as it:
                for entry in sorted(it, key=lambda e: e.name):
                    child_item = QTreeWidgetItem()
                    child_item.setText(0, entry.name)
                    child_item.setText(1, "待统计")
                    child_item.setText(2, "待统计")
                    child_item.setText(3, "待统计")
                    child_item.setText(4, "未统计")
                    if entry.is_dir():
                        child_item.setIcon(0, self.style().standardIcon(self.style().StandardPixmap.SP_DirIcon))
                        child_item.setChildIndicatorPolicy(QTreeWidgetItem.ChildIndicatorPolicy.ShowIndicator)
                    else:
                        child_item.setIcon(0, self.style().standardIcon(self.style().StandardPixmap.SP_FileIcon))
                    item.addChild(child_item)
        except Exception as e:
            print(f"Error expanding folder {path}: {str(e)}")

    def get_full_path(self, item):
        """获取树形项目对应的完整路径（支持多层级）"""
        path_parts = []
        current = item
        while current:
            path_parts.insert(0, current.text(0))
            current = current.parent()
        # 在选中的文件中查找匹配的路径
        for selected_path in self.selected_files:
            if os.path.basename(selected_path) == path_parts[0]:
                if len(path_parts) == 1:
                    return selected_path
                # 拼接子路径
                candidate = os.path.join(selected_path, *path_parts[1:])
                if os.path.exists(candidate):
                    return candidate
        return None 